using System;
using System.Collections;
using System.IO;
using System.Text;
using System.Diagnostics;

namespace A4G.Core
{
	public class GibbsSampler
	{
		private Random _rnd = new Random();

		private const double _pA = 0.25;
		private const double _pC = 0.25;
		private const double _pG = 0.25;
		private const double _pT = 0.25;

		private const int _aA = 1;
		private const int _aC = 1;
		private const int _aG = 1;
		private const int _aT = 1;
		private const int _sum_a = _aA + _aC + _aG + _aT;

		private object _A = 'A';
		private object _C = 'C';
		private object _G = 'G';
		private object _T = 'T';

		public GibbsSampler(Sequence[] fastaSequences, int motifLength, string filename)
		{
			_fastaSequences = fastaSequences;

			ArrayList sequences = new ArrayList();
			foreach (Sequence sequence in fastaSequences)
			{
				sequences.Add(sequence.SequenceString);
			}
			_sequences = (string[])sequences.ToArray(typeof(string));

			_motifLength = motifLength;
			_positions = new int[fastaSequences.Length];
			_filename = filename;

			_positionsWithBestScore = new int[Sequences.Length];

			InitPositions();
			Init();
		}

		private readonly Sequence[] _fastaSequences;
		public Sequence[] FastaSequences
		{
			get
			{
				return _fastaSequences;
			}
		}

		private readonly string[] _sequences;
		public string[] Sequences
		{
			get
			{
				return _sequences;
			}
		}

		private readonly int _motifLength = 0;
		public int MotifLength
		{
			get
			{
				return _motifLength;
			}
		}

		private readonly int[] _positions;
		public int[] Positions
		{
			get
			{
				return _positions;
			}
		}

		private readonly string _filename;
		public string Filename
		{
			get
			{
				return _filename;
			}
		}

		private readonly int[] _positionsWithBestScore;
		public int[] PositionsWithBestScore
		{
			get
			{
				return _positionsWithBestScore;
			}
		}

		private int _bestScore = 0;
		public int BestScore
		{
			get
			{
				return _bestScore;
			}
			set
			{
				_bestScore = value;
			}
		}


		private readonly ArrayList _positionFrequencies = new ArrayList();
		public ArrayList PositionFrequencies
		{
			get
			{
				return _positionFrequencies;
			}
		}

		public void Start(int maxCycles)
		{
			string positionFilename = Path.ChangeExtension(Filename, "positions");
			StreamWriter writer = new StreamWriter(positionFilename);

			try
			{
				PositionFrequencies.Clear();
				for (int i = 0; i < Sequences.Length; i++)
				{
					PositionFrequencies.Add(new Hashtable());
				}

				BestScore = 0;

				int displayCount = 0;
				int count = 0;
				while (count++ < maxCycles)
				{
					// Randomly select the sequence to omit
					int t = _rnd.Next(Sequences.Length);

					// Remove the selected sequence from the frequency matrix
					CalculateFrequencyMatrix(t, -1);

					// Get the new index in the omitted sequence
					Positions[t] = GetNewIndex(Sequences[t]);

					// Add the omitted sequence to the frequency matrix
					CalculateFrequencyMatrix(t, 1);

					int tmpScore = GetScore();
					if (tmpScore > BestScore)
					{
						BestScore = tmpScore;
						Positions.CopyTo(PositionsWithBestScore, 0);
					}

					int tmp = count / 1000;
					if (tmp > displayCount)
					{
						displayCount = tmp;
						Console.WriteLine(displayCount * 1000);
					}

					foreach (int position in Positions)
					{
						writer.Write(string.Format("{0}\t", position));
						//Debug.Write(string.Format("{0}\t", position));
					}
					writer.WriteLine("");
					//Debug.WriteLine("");
				}
			}
			finally
			{
				writer.Close();
			}
		}

		public void WriteMotifs(string filename, int W, int cycles)
		{
			StreamWriter outputWriter = new StreamWriter(filename);// Path.ChangeExtension(filename, "a4g.output"));
			try
			{
				outputWriter.WriteLine("Gibbs sampling algorithm.");
				outputWriter.WriteLine("Input file: " + filename);
				outputWriter.WriteLine("W: " + W);
				outputWriter.WriteLine("Cycles: " + cycles);
				outputWriter.WriteLine("");

				int count = 1;
				outputWriter.WriteLine("Sequences:");
				foreach (Sequence seq in FastaSequences)
				{
					outputWriter.WriteLine(string.Format("#{0} {1}", count++, seq.Name));
				}
				outputWriter.WriteLine("");

				outputWriter.WriteLine("Last Motifs:");
				string[] motifs = GetMotifs(Positions);
				count = 0;
				foreach (string motif in motifs)
				{
					outputWriter.WriteLine(string.Format("{0}\t{1}\t{2}", count+1, motif, Positions[count++]));
				}
				outputWriter.WriteLine("");

				outputWriter.WriteLine(string.Format("Best Score Motifs (score was {0}):", BestScore));
				motifs = GetMotifs(PositionsWithBestScore);
				count = 0;
				foreach (string motif in motifs)
				{
					outputWriter.WriteLine(string.Format("{0}\t{1}\t{2}", count+1, motif, PositionsWithBestScore[count++]));
				}
			}
			finally
			{
				outputWriter.Close();
			}
		}

		private string[] GetMotifs(int[] positions)
		{
			ArrayList motifs = new ArrayList();
			for (int k = 0; k < positions.Length; k++)
			{
				motifs.Add(Sequences[k].Substring(positions[k], MotifLength));
			}
			return (string[]) motifs.ToArray(typeof(string));
		}

		private int GetScore()
		{
			int score = 0;
			int[] tmp = new int[4];
			for (int r = 0; r < MotifLength; r++)
			{
				int[] frequencies = (int[]) F[_A];
				tmp[0] = frequencies[r];

				frequencies = (int[]) F[_C];
				tmp[1] = frequencies[r];

				frequencies = (int[]) F[_G];
				tmp[2] = frequencies[r];

				frequencies = (int[]) F[_T];
				tmp[3] = frequencies[r];

				score += GibbsSampler.GetMax(tmp);
			}
			return score;
		}

		private static int GetMax(int[] numbers)
		{
			int max = int.MinValue;
			for (int i = 0; i < numbers.Length; i++)
			{
				max = Math.Max(numbers[i], max);
			}
			return max;
		}

		private readonly Hashtable _a = new Hashtable();
		public Hashtable a
		{
			get
			{
				return _a;
			}
		}

		private readonly Hashtable _p0 = new Hashtable();
		public Hashtable P0
		{
			get
			{
				return _p0;
			}
		}

		private readonly Hashtable _f = new Hashtable();
		public Hashtable F
		{
			get
			{
				return _f;
			}
		}

		private void Init()
		{
			a.Add(_A, _aA);
			a.Add(_C, _aC);
			a.Add(_G, _aG);
			a.Add(_T, _aT);

			P0.Add(_A, _pA);
			P0.Add(_C, _pC);
			P0.Add(_G, _pG);
			P0.Add(_T, _pT);

			F.Add(_A, new int[MotifLength]);
			F.Add(_C, new int[MotifLength]);
			F.Add(_G, new int[MotifLength]);
			F.Add(_T, new int[MotifLength]);
			CalculateFrequencyMatrix();
		}

		private void CalculateFrequencyMatrix()
		{
			for (int k = 0; k < Sequences.Length; k++)
			{
				CalculateFrequencyMatrix(k, 1);
			}
		}

		private void CalculateFrequencyMatrix(int k, int increment)
		{
			string sequence = Sequences[k];
			for (int r = 0; r < MotifLength; r++)
			{
				char s = sequence[Positions[k] + r];
				int[] frequencies = (int[]) F[s];
				frequencies[r] += increment;
			}
		}

		private double CalculatePseudocountedProbability(char s, int frequency)
		{
			return ((double)(frequency + (int) a[s]) / ((Sequences.Length - 1) + _sum_a));
		}

		private int GetNewIndex(string sequence)
		{
			int rLength = sequence.Length - MotifLength;

			double sumQrOverPr = 0;
			double[] C = new double[rLength];
			for (int i = 0; i < rLength; i++)
			{
				double Qr = 1;
				double Pr = 1;
				for (int r = 0; r < MotifLength; r++)
				{
					char s = sequence[i + r];
					int[] frequencies = (int[]) F[s];
					Qr *= CalculatePseudocountedProbability(s, frequencies[r]);
					Pr *= (double) P0[s];
				}
				C[i] = (Qr / Pr);
				sumQrOverPr += C[i];
			}

			double tmp = 0;
			for (int r = 0; r < rLength; r++)
			{
				C[r] /= sumQrOverPr;
				C[r] += tmp;
				tmp = C[r];
			}

			double hit = _rnd.NextDouble() * tmp;

			int lower = 0;
			int upper = rLength;

			bool found = false;
			while (!found)
			{
				int middle = (upper + lower) / 2;
				if (C[middle] > hit)
				{
					upper = middle;
				}
				else
				{
					lower = middle;
				}

				if ((upper - lower) == 1)
				{
					found = true;
				}
			}

			return lower;
		}

		private void InitPositions()
		{
			for (int i = 0; i < Sequences.Length; i++)
			{
				Positions[i] = _rnd.Next(Sequences[i].Length - MotifLength);
			}
		}
	}
}
